package com.sunzm.spark.streaming

import org.apache.commons.lang3.StringUtils
import org.apache.commons.lang3.time.DateFormatUtils
import org.apache.spark.SparkConf
import org.apache.spark.streaming.{Durations, StreamingContext}
import org.apache.spark.streaming.dstream.DStream

object UpdateStateByKeyDemo {

  private val checkpointDirectory = "data/spark/streaming/ck/20210612001/"

  // Function to create and setup a new StreamingContext
  def functionToCreateContext(): StreamingContext = {
    val conf = new SparkConf()
      .setAppName(s"${this.getClass.getSimpleName.stripSuffix("$")}")
      .setMaster("local[*]")

    val ssc = new StreamingContext(conf, Durations.seconds(5))

    processFunc(ssc)

    //设置checkpoint目录
    ssc.checkpoint(checkpointDirectory)
    ssc
  }

  def main(args: Array[String]): Unit = {

    // Get StreamingContext from checkpoint data or create a new one
    val ssc = StreamingContext.getOrCreate(checkpointDirectory, functionToCreateContext _)

    ssc.start()
    ssc.awaitTermination()

  }

  def processFunc(ssc: StreamingContext) = {
    val lines = ssc.socketTextStream("82.156.210.70", 9999)

    val pairsRDD: DStream[(String, Int)] = lines.flatMap(_.split(","))
      .filter(line => {
        //只要不包含hello的
        !StringUtils.contains(line, "hello")
      })
      .map((_, 1))

    val wordCountDStream: DStream[(String, Int)] = pairsRDD.updateStateByKey(updateFunction)

    wordCountDStream.foreachRDD((rdd, time) => {

      val milliseconds = time.milliseconds

      val timeStr = DateFormatUtils.format(milliseconds, "yyyy-MM-dd HH:mm:ss")

      if (!rdd.isEmpty()) {

        println(s"当前计算时间: ${timeStr}")

        rdd.foreach(line => {
          println(s"${line}")
        })
      } else {
        println(s"没有需要计算的数据: ${timeStr}")
      }

    })
  }

  /**
   *
   * @param newValues  本批次新收到的数据集合(前5秒收到的数据)
   * @param stateCount 历史状态中保存的次数
   * @return
   */
  def updateFunction(newValues: Seq[Int], stateCount: Option[Int]): Option[Int] = {
    //本次收到的数据条数
    val newCount = newValues.sum

    //历史的数据条数
    val oldCount = stateCount.getOrElse(0)

    val sum = newCount + oldCount

    Some(sum)
  }
}
